import pandas as pd
import matplotlib.pyplot as plt
import os

PATH = r'C:\Users\Jeff Levy\Desktop\replication'

def generate_composite_plot(df):
    suptitle = 'Makeup of EPU Composite'
    ntl = df.groupby('date').mean() #equal weight ntl aggregate
    ntl = ntl.resample('Q').mean()

    epuc = ntl['EPU_Composite']
    epun = ntl['EPU_National']
    epus = ntl['EPU_State']
    x = ntl.index

    #set up stacked bar
    middle = (epun + epus) - epuc
    top = epus - middle
    bottom = epun - middle

    st_color, nt_color = 'forestgreen', 'orange'
    alpha = 1

    plt.rcParams.update({'hatch.color': st_color})
    width = 70

    fig, ax = plt.subplots(1, figsize=(12,6))
    ax.bar(x, top, width, bottom=bottom+middle, color=st_color, label='State (EPU-S)', alpha=alpha)
    ax.bar(x, middle, width, bottom=bottom, color=nt_color, hatch='//////', label='EPU-N and EPU-S Overlap', alpha=alpha)
    ax.bar(x, bottom, width, color=nt_color, label='National (EPU-N)', alpha=alpha)
    ax.legend(loc='upper center')

    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.margins(0.01)

    fig.supylabel('Mean Quarterly EPU-C', fontsize='xx-large', x=0.06)
    print('Saving: ', os.path.join(PATH, 'figure_4.png'))
    fig.savefig(os.path.join(PATH, 'figure_4.png'))
    plt.close()

if __name__ == '__main__':
    df = pd.read_csv(os.path.join(PATH, 'data.csv'), parse_dates=True, index_col=['state', 'date'])
    generate_composite_plot(df)
